﻿using ARSoft.Tools.Net;
using ARSoft.Tools.Net.Dns;
using System;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Threading.Tasks;

namespace DSVISUAL
{
    /// <summary>
    /// DNS代理服务器
    /// </summary>
    public class DnsProxy
    {
        public event EventHandler<MyQueryReceivedEventArgs> QueryReceived;
        private static DnsProxy _instance;
        public static DnsProxy Instance
        {
            get
            {
                if (_instance == null)
                {
                    _instance = new DnsProxy();
                }
                return _instance;
            }
        }

        public DnsMessage QueryFromSOA(IPAddress[] addresses, DnsMessage query)
        {
            var client = new DnsClient(addresses, 10000) { IsResponseValidationEnabled = true };
            var response = client.SendMessage(query);
            return response;
        }
        public ProxyMode Mode { get; set; }
        private bool flagRunning = false;
        private IPAddress localhost = IPAddress.Parse("127.0.0.1");
        private IPAddress DNS1 = IPAddress.Parse("61.139.2.69");
        private IPAddress DNS2 = IPAddress.Parse("61.139.2.69");
        public event EventHandler StateChanged;
        private DnsProxy()
        {
        }

        public void StopProxy()
        {
            try
            {
                flagRunning = false;
                Server.QueryReceived -= Server_QueryReceived;
                Server.Stop();
                StateChanged?.Invoke(this, EventArgs.Empty);
            }
            catch
            { }
        }
        DnsServer Server = new DnsServer(100, 100);

        public void StartProxy()
        {
            if (flagRunning)
                return;

            flagRunning = true;
            StateChanged?.Invoke(this, EventArgs.Empty);

            Server.QueryReceived += Server_QueryReceived;
            Server.Start();
        }

        private Task Server_QueryReceived(object sender, QueryReceivedEventArgs eventArgs)
        {
            return Task.Factory.StartNew((state) =>
           {
               var e = state as QueryReceivedEventArgs;
               e.Response = ((DnsMessage)e.Query).CreateResponseInstance();
               ProcessDNSQuery(e);
               return e;
           }, eventArgs);
        }

        class QueryData
        {
            public QueryData(byte[] data, EndPoint client)
            {
                Data = data;
                EndPoint = client;
            }

            public byte[] Data { get; set; }
            public EndPoint EndPoint { get; set; }
        }

        public enum ProxyMode
        {
            LocalFirst,
            PublicFirst,
            LocalOnly,
            PublicOnly,
        }

        private void QueryFromPublicDNS(QueryReceivedEventArgs e)
        {
            var query = (DnsMessage)e.Query;
            var response = (DnsMessage)e.Response;
            DnsClient client = new DnsClient(new IPAddress[] { DNS1, DNS2 }, 5000) { IsResponseValidationEnabled = true };
            var resp = client.SendMessage(query);
            e.Response = resp ?? e.Response;
        }

        private void QueryFromPrivateDNS(QueryReceivedEventArgs dns)
        {
            DomainManager.Instance.Resolute(dns);
            var queryLoop = 5;
            while (queryLoop-- > 0)
            {
                if ((dns.Response as DnsMessage).AnswerRecords.Any(x => x.RecordType == RecordType.A || x.RecordType == RecordType.Aaaa || x.RecordType == RecordType.Soa))
                {
                    break;
                }
                var cnameList = (dns.Response as DnsMessage).AnswerRecords.Where(x => x.RecordType == RecordType.CName).ToList();
                if (!cnameList.Any())
                {
                    break;
                }
                foreach (CNameRecord item in cnameList)
                {
                    var q = (dns.Query as DnsMessage).Questions.Find(x => x.Name == item.Name);
                    (dns.Query as DnsMessage).Questions.Remove(q);
                    (dns.Response as DnsMessage).AnswerRecords.Remove(item);
                    (dns.Query as DnsMessage).Questions.Add(new DnsQuestion(item.CanonicalName, RecordType.A, RecordClass.INet));
                    DomainManager.Instance.Resolute(dns);
                }
            }
            dns.Response.responderAddress = IPAddress.Loopback;
        }


        private void OnDNSQuery(QueryReceivedEventArgs args, IPAddress server, TimeSpan elapsed)
        {
            QueryReceived?.Invoke(this, new MyQueryReceivedEventArgs { EventArgs = args, ServerEndPoint = new IPEndPoint(server, 53), TimeElapsed = elapsed });
        }
        private bool debug;
        private void ProcessDNSQuery(QueryReceivedEventArgs e)
        {

            var stopwatch1 = Stopwatch.StartNew();

            var query = (DnsMessage)e.Query;
            var response = (DnsMessage)e.Response;
            if (query.Questions.Count == 1 && query.Questions[0].RecordType == RecordType.Ptr)
            {
                if (query.Questions[0].Name.ToString() == "1.0.0.127.in-addr.arpa." || query.Questions[0].Name.ToString() == "1.0.0.127.in-addr.arpa")
                {
                    response.AnswerRecords.Add(new PtrRecord(query.Questions[0].Name, 1440, DomainName.Parse(Environment.MachineName)));
                    response.responderAddress = IPAddress.Loopback;
                }
            }
            else if (query.Questions.Count == 1 && (query.Questions[0].RecordType == RecordType.A || query.Questions[0].RecordType == RecordType.Aaaa))
            {
                //if (debug)
                //{
                //    return;
                //}
                //debug = true;
                if (query.Questions[0].Name.ToString().Equals(Environment.MachineName + ".", StringComparison.OrdinalIgnoreCase))
                {
                    response.AnswerRecords.Add(new ARecord(query.Questions[0].Name, 1440, e.RemoteEndpoint.Address));
                }
                else if (query.Questions[0].Name.ToString().EndsWith(".lan."))
                {
                    QueryFromPrivateDNS(e);
                    if (!(e.Response as DnsMessage).AnswerRecords.Any())
                    {
                        var entry = Dns.GetHostEntry(query.Questions[0].Name.ToString().Replace(".lan.", ""));
                        foreach (var item in entry.AddressList)
                        {
                            response.AnswerRecords.Add(new ARecord(query.Questions[0].Name, 1440, item));
                        }
                        if (!entry.AddressList.Any())
                        {
                            ((DnsMessage)e.Response).ReturnCode = ReturnCode.NxDomain;
                        }
                        response.responderAddress = IPAddress.Loopback;
                    }
                }
                else
                {
                    ProcessQuery(e);
                }
            }
            else
            {
                ProcessQuery(e);
            }


            response = (DnsMessage)e.Response; // must reasign
            //relay the SOA.
            SoaRecord soaRecord = response.AuthorityRecords.OfType<SoaRecord>().FirstOrDefault();

            int soaTreeLevelMax = 5;
            int soaTreeLevel = 0;
            while (soaTreeLevel < soaTreeLevelMax && !response.AnswerRecords.Any() && soaRecord != null)
            {
                var soaQuery = DnsMessage.CreateQuery(soaRecord.Name,
                    query.Questions[0].RecordType,
                    query.Questions[0].RecordClass);
                soaQuery.IsEDnsEnabled = false;


                var mResolveQuery = DnsMessage.CreateQuery(soaRecord.MasterName);
                mResolveQuery.IsEDnsEnabled = false;
                var mResolve = this.QueryFromSOA(new IPAddress[] { IPAddress.Loopback }, mResolveQuery);
                if (mResolve == null || mResolve.ReturnCode != ReturnCode.NoError)
                    break;

                var addressList = mResolve.AnswerRecords.Where(x => x.RecordType == RecordType.A || x.RecordType == RecordType.Aaaa).ToArray();
                if (!addressList.Any())
                {
                    break;
                }
                var st2 = Stopwatch.StartNew();
                var soaResponse = this.QueryFromSOA(addressList.Select(x => (x as AddressRecordBase).Address).ToArray(), soaQuery);
                if (soaResponse == null)
                {
                    break;
                }
                //insert the response to ui listview
                this.OnDNSQuery(new QueryReceivedEventArgs(soaQuery, ProtocolType.Tcp, e.LocalEndPoint, null) { Response = soaResponse }, soaResponse.responderAddress, st2.Elapsed);

                if (soaResponse.AnswerRecords.Any(x => x.RecordType == RecordType.A || x.RecordType == RecordType.Aaaa))
                {
                    //授权服务器的A记录
                    response.AdditionalRecords.Add(new ARecord(soaRecord.MasterName, 120, soaResponse.responderAddress));

                    //授权服务器的NS记录
                    response.AuthorityRecords.Add(new NsRecord(soaRecord.Name, 120, soaRecord.MasterName));

                    //从授权服务器查询得到的Answer记录
                    response.AnswerRecords.AddRange(soaResponse.AnswerRecords);
                    response.IsAuthoritiveAnswer = true;
                    break;
                }
                else if (soaResponse.AuthorityRecords.OfType<SoaRecord>().Any())
                {
                    soaRecord = soaResponse.AuthorityRecords.OfType<SoaRecord>().FirstOrDefault();
                    //m.AuthorityRecords.AddRange(soaResponse.AuthorityRecords);
                    soaTreeLevel++;
                }
                else if (soaResponse.AnswerRecords.Any(x => x.RecordType == RecordType.Opt))
                {
                    response.AnswerRecords.AddRange(soaResponse.AnswerRecords);
                    break;
                }
                else
                {
                    break;
                }
            }
            try
            {
                if (response.AdditionalRecords.Count + response.AnswerRecords.Count + response.AuthorityRecords.Count == 0)
                {
                    response.ReturnCode = ReturnCode.NxDomain;
                }
                this.OnDNSQuery(e, e.Response.responderAddress, stopwatch1.Elapsed);
            }
            catch { }
        }

        public void SetDNS1(IPAddress dns2)
        {
            this.DNS1 = dns2;
        }
        public void SetDNS2(IPAddress dns2)
        {
            this.DNS2 = dns2;
        }



        public bool IsProxyRuning
        {
            get
            {
                return flagRunning;
            }
        }

        public void ProcessQuery(QueryReceivedEventArgs e)
        {
            if (Mode == ProxyMode.PublicOnly)
            {
                QueryFromPublicDNS(e);
            }
            else if (Mode == ProxyMode.LocalOnly)
            {
                QueryFromPrivateDNS(e);
            }
            else if (Mode == ProxyMode.LocalFirst)
            {
                QueryFromPrivateDNS(e);
                if (!(e.Response as DnsMessage).AnswerRecords.Any() && !(e.Response as DnsMessage).AuthorityRecords.Any() && !(e.Response as DnsMessage).AdditionalRecords.Any())
                {
                    QueryFromPublicDNS(e);
                }
            }
            else if (Mode == ProxyMode.PublicFirst)
            {
                QueryFromPublicDNS(e);
                if (!(e.Response as DnsMessage).AnswerRecords.Any() && !(e.Response as DnsMessage).AuthorityRecords.Any() && !(e.Response as DnsMessage).AdditionalRecords.Any())
                {
                    QueryFromPrivateDNS(e);
                }
            }
        }
    }
}

